import rospy

from webots_ros.srv import set_float
from webots_ros.srv import set_int
from webots_ros.srv import get_float
from webots_ros.msg import Float64Stamped
from webots_ros.msg import Int32Stamped

class ChassisManager():
    def __init__(self):
        webots_ns = 'robot'
        joint1_client = '{}/small_wheel1/{}'
        joint2_client = '{}/small_wheel2/{}'
        joint3_client = '{}/small_wheel3/{}'

        joint4_client = '{}/wheel1/{}'
        joint5_client = '{}/wheel2/{}'
        joint6_client = '{}/wheel3/{}'
        
        self.angle = 1.5708
        self.position = 99999
        self.scale = 10/310.0
        self.wheel_angle = 0
        self.radius = 0.1
        # create proxy methods to set motor positions
        rospy.wait_for_service(joint1_client.format(webots_ns, 'set_position'))
        self.j1_set_position = rospy.ServiceProxy(
            joint1_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint2_client.format(webots_ns, 'set_position'))
        self.j2_set_position = rospy.ServiceProxy(
            joint2_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint3_client.format(webots_ns, 'set_position'))
        self.j3_set_position = rospy.ServiceProxy(
            joint3_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint4_client.format(webots_ns, 'set_position'))
        self.j4_set_position = rospy.ServiceProxy(
            joint4_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint5_client.format(webots_ns, 'set_position'))
        self.j5_set_position = rospy.ServiceProxy(
            joint5_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint6_client.format(webots_ns, 'set_position'))
        self.j6_set_position = rospy.ServiceProxy(
            joint6_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        # create proxy methods to set motor velocities
        rospy.wait_for_service(joint1_client.format(webots_ns, 'set_velocity'))
        self.j1_set_velocity = rospy.ServiceProxy(
            joint1_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint2_client.format(webots_ns, 'set_velocity'))
        self.j2_set_velocity = rospy.ServiceProxy(
            joint2_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint3_client.format(webots_ns, 'set_velocity'))
        self.j3_set_velocity = rospy.ServiceProxy(
            joint3_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint4_client.format(webots_ns, 'set_velocity'))
        self.j4_set_velocity = rospy.ServiceProxy(
            joint4_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint5_client.format(webots_ns, 'set_velocity'))
        self.j5_set_velocity = rospy.ServiceProxy(
            joint5_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint6_client.format(webots_ns, 'set_velocity'))
        self.j6_set_velocity = rospy.ServiceProxy(
            joint6_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service("/robot/small_wheel1_sensor/enable")
        self.small_wheel1_sensor_enable = rospy.ServiceProxy("/robot/small_wheel1_sensor/enable", set_int, persistent=True)
        self.small_wheel1_sensor_enable(1)
        rospy.Subscriber('/robot/small_wheel1_sensor/value', Float64Stamped, self.WSCallback)

        rospy.wait_for_service("/robot/keyboard/enable")
        self.keyborad_enable = rospy.ServiceProxy("/robot/keyboard/enable", set_int, persistent=True)
   
    def WSCallback(self,WS_msg):
        self.wheel_angle = WS_msg.data

    def KeyBoardCallback(self,key_value_msg):
        #forward
        if(key_value_msg.data==315):
            # self.position = self.position+self.scale
            self.set_velocity(self.position,0.2)
        #back
        if(key_value_msg.data==317):
            # self.position = self.position-self.scale
            self.set_velocity(-self.position,0.2)
        #right
        if(key_value_msg.data==316):
            # self.angle = self.angle+self.scale
            self.set_angle(self.angle,0.2)
        #left
        if(key_value_msg.data==314):
            # self.angle = self.angle-self.scale
            self.set_angle(-self.angle,0.2)
        #stop
        if(key_value_msg.data==32):
            # self.angle = self.angle-self.scale
            self.set_velocity(self.position,0)
            self.set_angle(self.angle,0)
# rosrun topic_tools throttle messages /robot/keyboard/key 1.0 /newkey


    def set_velocity(self, joint_position , joint_speed = 0.1):
        self.j1_set_velocity(joint_speed)
        self.j2_set_velocity(joint_speed)
        self.j3_set_velocity(joint_speed)

        self.j1_set_position(joint_position)
        self.j2_set_position(joint_position)
        self.j3_set_position(joint_position)

    def set_position(self, joint_position , joint_speed = 0.1):
        print(self.wheel_angle,joint_position/self.radius)
        target = self.wheel_angle + joint_position/self.radius
        while(abs(self.wheel_angle - target)>0.05):
            self.j1_set_velocity(joint_speed)
            self.j2_set_velocity(joint_speed)
            self.j3_set_velocity(joint_speed)

            self.j1_set_position(target)
            self.j2_set_position(target)
            self.j3_set_position(target)
            rospy.sleep(0.1)

    def set_angle(self, joint_position, joint_speed = 0.1):  
        self.j4_set_velocity(joint_speed)
        self.j5_set_velocity(joint_speed)
        self.j6_set_velocity(joint_speed)
        self.j4_set_position(joint_position)
        self.j5_set_position(joint_position)     
        self.j6_set_position(joint_position)  

    def run(self):
        self.keyborad_enable(1)
        rospy.Subscriber('/newkey', Int32Stamped, self.KeyBoardCallback)
        print("chassis manager initialization is done")         
        rospy.spin()

if __name__ == "__main__":

    rospy.init_node('ChassisManager', anonymous= False)

    cm = ChassisManager()
    # cm.run()
    cm.set_angle(1.5708,0.05)
    # cm.set_position(0.5,0.1)


    # while rospy.is_shutdown():
    #     print("kill the node")
    # cm.set_velocity()

